import requests
import json
import subprocess
import time
from datetime import datetime

def get_available_models():
    """获取本地可用的Ollama模型列表"""
    try:
        result = subprocess.run(['ollama', 'list'], capture_output=True, text=True, check=True)
        models = []
        # 解析模型列表，跳过标题行
        for line in result.stdout.split('\n')[1:]:
            if line.strip():
                model_name = line.split()[0]
                models.append(model_name)
        return models
    except subprocess.CalledProcessError as e:
        print(f"获取模型列表失败: {e}")
        return []

def chat_with_ollama(model, prompt):
    """与指定模型进行对话"""
    url = "http://localhost:11434/api/chat"
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "stream": False,
        "options": {
            "temperature": 0.7,
            "num_ctx": 4096
        }
    }
    
    try:
        start_time = time.time()
        response = requests.post(url, json=payload)
        response_time = time.time() - start_time
        response.raise_for_status()
        return response.json(), response_time
    except requests.exceptions.RequestException as e:
        print(f"API请求错误: {e}")
        return None, 0

def main():
    print("=== 多模型对话测试脚本 ===")
    print("确保Ollama服务正在运行")
    
    # 获取可用模型列表
    available_models = get_available_models()
    if not available_models:
        return
    
    print("\n可用的模型列表:")
    for i, model in enumerate(available_models, 1):
        print(f"{i}. {model}")
    
    # 用户输入
    try:
        cycles = int(input("\n请输入每个模型的对话轮次: "))
        question = input("请输入要提问的问题: ")
        
        # 选择模型
        model_count = int(input("请输入要使用的模型数量: "))
        selected_models = []
        for _ in range(model_count):
            model_name = input(f"请输入第{_+1}个模型名称（完整名称）: ").strip()
            if model_name in available_models:
                selected_models.append(model_name)
            else:
                print(f"错误：模型 {model_name} 不存在，请重新输入")
                return
    except ValueError:
        print("输入无效，请确保输入正确的数字和模型名称")
        return
    
    # 遍历所有选定模型
    for model in selected_models:
        print(f"\n=== 开始测试模型 {model} ===")
        
        conversation_data = {
            "model": model,
            "question": question,
            "total_cycles": cycles,
            "conversations": [],
            "timing": {
                "total_time": 0,
                "average_time": 0
            }
        }
        
        # 执行对话循环
        successful_cycles = 0
        total_response_time = 0
        
        for i in range(cycles):
            print(f"正在进行第 {i+1}/{cycles} 轮对话...")
            
            response, response_time = chat_with_ollama(model, question)
            
            if response:
                successful_cycles += 1
                total_response_time += response_time
                
                # 记录对话详情
                conversation_data["conversations"].append({
                    "cycle": i+1,
                    "response": response.get("message", {}).get("content", "无响应"),
                    "response_time": round(response_time, 2),
                    "timestamp": datetime.now().isoformat(),
                    "context_tokens": len(response.get("context", [])),
                    "model_details": {
                        "model": response.get("model"),
                        "created_at": response.get("created_at"),
                        "total_duration": response.get("total_duration")
                    }
                })
                
                # 打印简略信息
                preview = response["message"]["content"][:50] + "..." if len(response["message"]["content"]) > 50 else response["message"]["content"]
                print(f"响应预览: {preview}")
                print(f"本次耗时: {response_time:.2f}s\n")
            else:
                conversation_data["conversations"].append({
                    "cycle": i+1,
                    "error": "请求失败"
                })
        
        # 计算时间统计
        if successful_cycles > 0:
            conversation_data["timing"]["total_time"] = round(total_response_time, 2)
            conversation_data["timing"]["average_time"] = round(total_response_time / successful_cycles, 2)
        
        # 生成文件名
        model_safe = model.replace(":", "-").replace(" ", "_")
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{model_safe}_cycles{cycles}_{timestamp}.json"
        
        # 保存文件
        try:
            with open(filename, "w", encoding="utf-8") as f:
                json.dump(conversation_data, f, ensure_ascii=False, indent=2)
            print(f"对话记录已保存到 {filename}")
            print(f"总耗时: {conversation_data['timing']['total_time']}s")
            print(f"平均响应时间: {conversation_data['timing']['average_time']}s\n")
        except Exception as e:
            print(f"保存文件时出错: {e}")

if __name__ == "__main__":
    main()